from __future__ import annotations

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.carla_scenario import CarlaScenario
from driving_gym.misc.roach.core.obs_manager.obs_manager_handler import (
    ObsManagerHandler,
)
from driving_gym.misc.roach.core.task_actor.ego_vehicle.ego_vehicle_handler import (
    EgoVehicleHandler,
    TaskVehicle,
)
from driving_gym.misc.roach.utils.traffic_light import TrafficLightHandler
from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter
from driving_gym.simulation.adapter.carla.carla_config import (
    CarlaActorConfig,
    carla_util,
)


class CarlaRoachObs(BaseObs):
    adapter: CarlaAdapter

    def __init__(self, config: dict, adapter: CarlaAdapter):
        super().__init__(config, adapter)

        obs_config = self.config.copy()
        for k in list(obs_config.keys()):
            if not isinstance(obs_config[k], dict):
                del obs_config[k]

        self._handler = ObsManagerHandler({self.actor_id: obs_config})
        self._ego_vehicle = None

    def reset(self):
        if TrafficLightHandler.carla_map is None:
            TrafficLightHandler.reset(self.adapter.world)

        self._handler.clean()
        self._ego_vehicle = TaskVehicle.VEHICLES[self.actor_id] = None
        snap_shot = self.adapter.world.get_snapshot()
        self._timestamp = {
            "step": 0,
            "frame": snap_shot.timestamp.frame,
            "relative_wall_time": 0.0,
            "wall_time": snap_shot.timestamp.platform_timestamp,
            "relative_simulation_time": 0.0,
            "simulation_time": snap_shot.timestamp.elapsed_seconds,
            "start_frame": snap_shot.timestamp.frame,
            "start_wall_time": snap_shot.timestamp.platform_timestamp,
            "start_simulation_time": snap_shot.timestamp.elapsed_seconds,
        }
        return super().reset()

    def get_obs(self, snapshot: Snapshot, scenario: CarlaScenario):
        if self._ego_vehicle is None:
            self._ego_vehicle = self._wrap_vehicle(scenario)
            self._handler.reset({self.actor_id: self._ego_vehicle})

        self._update_timestamp()
        self._ego_vehicle.tick(self._timestamp)
        return self._handler.get_observation(self._timestamp)[self.actor_id]

    def get_observation_space(self):
        return self._handler.observation_space[self.actor_id]

    def _wrap_vehicle(self, scenario: CarlaScenario):
        ego_vehicle = self.adapter.get_actor(self.actor_id)
        carla_map = self.adapter._map

        task_vehicle = TaskVehicle.VEHICLES.get(self.actor_id, None)
        if task_vehicle is None:
            spawn_transforms = EgoVehicleHandler._get_spawn_points(carla_map)
            task = scenario.get_task(self.actor_id)
            route = task.get("route", None) or [task["end"]]
            target_transforms = [
                carla_util.carla_tranform(
                    CarlaActorConfig.parse_location(loc, carla_map)
                )
                for loc in route
            ]

            task_vehicle = TaskVehicle(
                ego_vehicle, target_transforms, spawn_transforms, scenario.endless
            )
            TaskVehicle.VEHICLES[self.actor_id] = task_vehicle
        return task_vehicle

    def _update_timestamp(self):
        snap_shot = self.adapter.world.get_snapshot()
        self._timestamp["step"] = (
            snap_shot.timestamp.frame - self._timestamp["start_frame"]
        )
        self._timestamp["frame"] = snap_shot.timestamp.frame
        self._timestamp["wall_time"] = snap_shot.timestamp.platform_timestamp
        self._timestamp["relative_wall_time"] = (
            self._timestamp["wall_time"] - self._timestamp["start_wall_time"]
        )
        self._timestamp["simulation_time"] = snap_shot.timestamp.elapsed_seconds
        self._timestamp["relative_simulation_time"] = (
            self._timestamp["simulation_time"]
            - self._timestamp["start_simulation_time"]
        )
